package com.ai.module.mj.controller;

import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.RandomUtil;
import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.util.DateUtils;
import com.ai.common.BaseResp;
import com.ai.module.mj.config.propertites.YmlProperties;
import com.ai.module.mj.constant.enums.TaskAction;
import com.ai.module.mj.constant.enums.TranslateWay;
import com.ai.module.mj.constant.stringConstant.Constants;
import com.ai.module.mj.constant.stringConstant.ReturnCode;
import com.ai.module.mj.domain.Task;
import com.ai.module.mj.domain.vo.SubmitResultVO;
import com.ai.module.mj.exception.BannedPromptException;
import com.ai.module.mj.loadbalancer.DiscordLoadBalancer;
import com.ai.module.mj.service.HighTaskStoreService;
import com.ai.module.mj.service.TaskService;
import com.ai.module.mj.service.TranslateService;
import com.ai.module.mj.util.BannedPromptUtils;
import com.ai.module.mj.util.SnowFlakeUtil;
import com.ai.module.mj.wss.handle.MessageHandler;
import com.alibaba.fastjson2.JSONObject;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import io.swagger.annotations.ApiParam;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author: MrHe
 * @create: 2023-12-12 01:00
 * @description:
 **/
@Slf4j
@RestController
@RequestMapping("/draw/test")
@RequiredArgsConstructor
@Api(hidden = true)
public class DrawTestController {

    private final TranslateService translateService;
    private final HighTaskStoreService taskStoreService;
    private final YmlProperties properties;
    private final TaskService taskService;

    @ApiOperation(value = "提交Imagine任务", hidden = true)
    @PostMapping("/image")
    public BaseResp imagine(@RequestBody JSONObject json) {
        String prompt = json.getString("prompt");
        if (prompt!= null){
            System.out.println("prompt = " + prompt);
            return  BaseResp.success("123290809893");
        }
        if (CharSequenceUtil.isBlank(prompt)) {
            return BaseResp.fail("prompt不能为空",ReturnCode.VALIDATION_ERROR);
        }

        prompt = prompt.trim();
        Task task = new Task();
        task.setId(System.currentTimeMillis() + RandomUtil.randomNumbers(3));
        task.setSubmitTime(DateUtils.getCurrentTime());
        String notifyHook = this.properties.getNotifyHook()  ;
        task.setProperty(Constants.TASK_PROPERTY_NOTIFY_HOOK, notifyHook);
        task.setProperty(Constants.TASK_PROPERTY_NONCE, SnowFlakeUtil.INSTANCE.nextId());
        task.setUserId(ThreadLocalContext.getUserId());
        task.setUserIntId(ThreadLocalContext.getUserLoginInfo().getId());
        task.setAction(TaskAction.IMAGINE);
        task.setPrompt(prompt);
        String promptEn = translatePrompt(prompt);
        try {
            BannedPromptUtils.checkBanned(promptEn);
        } catch (BannedPromptException e) {
            return BaseResp.fail( "可能包含敏感词",ReturnCode.BANNED_PROMPT)
                 ;
                    //.setProperty("promptEn", promptEn).setProperty("bannedWord", e.getMessage());
        }catch (Exception ex){
            ex.printStackTrace();
        }

        task.setPromptEn(promptEn);
        task.setDescription("/imagine " + prompt);
        SubmitResultVO submitResultVO = this.taskService.submitImagine(task);
        log.info("AI 绘画 返回值为[{}]",submitResultVO);

        return BaseResp.success(submitResultVO.getCode(),submitResultVO.getMessage(),submitResultVO.getTaskId());
    }
    private String translatePrompt(String prompt) {
        if (TranslateWay.NULL.equals(this.properties.getTranslateWay()) || CharSequenceUtil.isBlank(prompt)) {
            return prompt;
        }
        List<String> imageUrls = new ArrayList<>();
        Matcher imageMatcher = Pattern.compile("https?://[a-z0-9-_:@&?=+,.!/~*'%$]+\\x20+", Pattern.CASE_INSENSITIVE).matcher(prompt);
        while (imageMatcher.find()) {
            imageUrls.add(imageMatcher.group(0));
        }
        String paramStr = "";
        Matcher paramMatcher = Pattern.compile("\\x20+-{1,2}[a-z]+.*$", Pattern.CASE_INSENSITIVE).matcher(prompt);
        if (paramMatcher.find()) {
            paramStr = paramMatcher.group(0);
        }
        String imageStr = CharSequenceUtil.join("", imageUrls);
        String text = prompt.substring(imageStr.length(), prompt.length() - paramStr.length());
        if (CharSequenceUtil.isNotBlank(text)) {
            text = this.translateService.translateToEnglish(text).trim();
        }
        return imageStr + text + paramStr;
    }

    private final DiscordLoadBalancer discordLoadBalancer;

    @ApiOperation(value = "指定ID获取任务", hidden = true)
    @GetMapping("/{id}/fetch")
    public Task fetch(@ApiParam(value = "任务ID") @PathVariable String id) {
        Task task = this.taskStoreService.get(id);
        if (task==null){
            return MessageHandler.storeTaskMap.get(id);
        }
        return task;
    }

}
